package top.wshape1.rpc.registry;

import com.alibaba.nacos.api.exception.NacosException;
import com.alibaba.nacos.api.naming.NamingFactory;
import com.alibaba.nacos.api.naming.NamingService;
import com.alibaba.nacos.api.naming.listener.Event;
import com.alibaba.nacos.api.naming.listener.EventListener;
import com.alibaba.nacos.api.naming.listener.NamingEvent;
import com.alibaba.nacos.api.naming.pojo.Instance;
import lombok.extern.slf4j.Slf4j;
import top.wshape1.rpc.common.constant.RoleType;
import top.wshape1.rpc.common.entity.ServiceMeta;
import top.wshape1.rpc.config.JJRpcConfig;

import java.util.Collections;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

/**
 * @author Wshape1
 * @version 1.0
 * @description
 * @since 2024-02-21
 */

@Slf4j
public class NacosRegistry implements RpcRegistry {

    private final static int RETRY_TIMES = 3;

    private final static int RETRY_INTERVAL = 1000;

    private NamingService namingService;

    private final String group;

    private final boolean enableNacosLoadBalance;

    /**
     * Service cache
     * ServiceName -> List<ServiceMeta>
     */
    private final ConcurrentHashMap<String, List<ServiceMeta>> SERVICE_CACHE = new ConcurrentHashMap<>();

    /**
     * Service listener cache
     * ServiceName -> EventListener
     */
    private final ConcurrentHashMap<String, EventListener> LISTENER_CACHE = new ConcurrentHashMap<>();

    public NacosRegistry() {
        String address = JJRpcConfig.CONFIG.getRegistry().getNacos().getAddress();
        for (int i = 0; i < RETRY_TIMES; i++) {
            try {
                this.namingService = NamingFactory.createNamingService(address);
                break;
            } catch (NacosException e) {
                log.error("init NacosRegistry failed, retry times:{}/{}, error: ", i + 1, RETRY_TIMES, e);
                try {
                    Thread.sleep(RETRY_INTERVAL);
                } catch (InterruptedException ignored) {
                }
            }
        }
        if (this.namingService == null)
            throw new RuntimeException("NacosRegistry init failed");
        this.group = JJRpcConfig.CONFIG.getRegistry().getNacos().getGroup();
        this.enableNacosLoadBalance = JJRpcConfig.CONFIG.getRegistry().getNacos().isEnableNacosLoadBalance();
    }

    @Override
    public void register(ServiceMeta serviceMeta) {
        try {
            Instance instance = new Instance();
            instance.setIp(serviceMeta.getIp());
            instance.setPort(serviceMeta.getPort());
            instance.setServiceName(serviceMeta.getServiceName());
            instance.setEphemeral(true);
            instance.setHealthy(true);
            instance.setEnabled(true);
            namingService.registerInstance(serviceMeta.getServiceName(), group, instance);
            if (!enableNacosLoadBalance)
                subscribe(serviceMeta.getServiceName());
        } catch (NacosException e) {
            log.error("register service failed, error: ", e);
            throw new RuntimeException(e);
        }
    }

    @Override
    public void unregister(ServiceMeta serviceMeta) {
        try {
            namingService.deregisterInstance(serviceMeta.getServiceName(), serviceMeta.getIp(), serviceMeta.getPort());
            if (!enableNacosLoadBalance)
                unsubscribe(serviceMeta.getServiceName());
        } catch (NacosException e) {
            log.error("unregister service failed, error: ", e);
            throw new RuntimeException(e);
        }
    }

    @Override
    public List<ServiceMeta> lookup(String serviceName) {
        List<ServiceMeta> list;
        if (enableNacosLoadBalance) {
            try {
                Instance instance = namingService.selectOneHealthyInstance(serviceName, group);
                ServiceMeta serviceMeta = new ServiceMeta();
                serviceMeta.setServiceName(serviceName);
                serviceMeta.setIp(instance.getIp());
                serviceMeta.setPort(instance.getPort());
                serviceMeta.setRole(RoleType.provider);
                list = Collections.singletonList(serviceMeta);
            } catch (NacosException e) {
                log.error("lookup service failed, error: ", e);
                throw new RuntimeException(e);
            }
        } else {
            list = SERVICE_CACHE.get(serviceName);
        }
        return list;
    }

    @Override
    public void subscribe(String serviceName) {
        EventListener listener = LISTENER_CACHE.get(serviceName);
        if (listener == null) {
            try {
                EventListener newListener = new ServiceEventListener();
                EventListener oldListener = LISTENER_CACHE.putIfAbsent(serviceName, newListener);
                if (oldListener == null) {
                    // 说明是第一次订阅
                    namingService.subscribe(serviceName, group, newListener);
                }
            } catch (NacosException e) {
                log.error("subscribe service failed, error: ", e);
                throw new RuntimeException(e);
            }
        }
    }

    @Override
    public void unsubscribe(String serviceName) {
        EventListener listener = LISTENER_CACHE.remove(serviceName);
        if (listener != null) {
            try {
                namingService.unsubscribe(serviceName, group, listener);
            } catch (NacosException e) {
                log.error("unsubscribe service failed, error: ", e);
                throw new RuntimeException(e);
            }
        }
    }

    @Override
    public void destroy() {
        LISTENER_CACHE.forEach((serviceName, listener) -> {
            try {
                namingService.unsubscribe(serviceName, group, listener);
            } catch (NacosException e) {
                log.error("unsubscribe service(name={}) failed, error: ", serviceName, e);
            }
        });
        LISTENER_CACHE.clear();
        SERVICE_CACHE.forEach((serviceName, list) -> {
            list.forEach(serviceMeta -> {
                try {
                    namingService.deregisterInstance(serviceName, serviceMeta.getIp(), serviceMeta.getPort());
                } catch (NacosException e) {
                    log.error("deregister service(name={}, address={}) failed, error: ", serviceName, serviceMeta.getAddress(), e);
                }
            });
        });
        SERVICE_CACHE.clear();
        try {
            namingService.shutDown();
        } catch (NacosException e) {
            log.error("destroy NacosRegistry failed, error: ", e);
        }
    }

    private class ServiceEventListener implements EventListener {

        @Override
        public void onEvent(Event event) {
            if (event instanceof NamingEvent) {
                NamingEvent namingEvent = (NamingEvent) event;
                List<ServiceMeta> list = namingEvent.getInstances()
                        .stream()
                        .map(instance -> {
                            ServiceMeta serviceMeta = new ServiceMeta();
                            serviceMeta.setServiceName(instance.getServiceName());
                            serviceMeta.setIp(instance.getIp());
                            serviceMeta.setPort(instance.getPort());
                            serviceMeta.setRole(RoleType.provider);
                            return serviceMeta;
                        }).collect(Collectors.toList());
                SERVICE_CACHE.put(namingEvent.getServiceName(), list);
            }
        }
    }

}
